from PyCmpltrtok.common import get_dir_name_ext
import os
import sys
import cv2 as cv
import numpy as np

NO_MEDIA_FILE_NAME = '.nomedia'


def check_dir_pic_exts(xdir, xexts):
    xset = set()
    xdict = {
        'n_dir': 0,
        'n_file': 0,
        'n_found': 0,
        'n_proc': 0,
    }
    _check_dir_pic_exts(xset, xdir, xdict, xexts)
    print()


def _check_dir_pic_exts(xset_ref, xdir, stat_ref, xexts):

    def print_stat():
        print(f'\rdirs={stat_ref["n_dir"]}, files={stat_ref["n_file"]}, found={stat_ref["n_found"]}, proc={stat_ref["n_proc"]}', end='', flush=True)

    stat_ref["n_dir"] += 1
    print_stat()
    xitems = os.listdir(xdir)
    xpaths = [os.path.join(xdir, xitem) for xitem in xitems if xitem != NO_MEDIA_FILE_NAME]
    for xpath in xpaths:
        if os.path.isfile(xpath):
            the_dir, the_base, xext = get_dir_name_ext(xpath)
            stat_ref["n_file"] += 1
            if xext in xexts:
                stat_ref['n_found'] += 1
                dest = os.path.join(the_dir, f'{the_base}.jpg')
                if os.path.exists(dest):
                    print(f'|{dest}| already exists!', file=sys.stderr, flush=True)
                else:
                    if 0:
                        # Cannot handle CJK path
                        img = cv.imread(xpath)
                        cv.imwrite(dest, img)
                    else:
                        if 1:
                            arr = np.fromfile(xpath, dtype=np.uint8)
                        else:
                            # >  - buf is not a numpy array, neither a scalar
                            # >  - Expected Ptr<cv::UMat> for argument 'buf'
                            with open(xpath, 'br') as f:
                                arr = f.read()
                        img = cv.imdecode(arr, cv.IMREAD_UNCHANGED)
                        ret, buf = cv.imencode('.jpg', img, [])
                        if not ret:
                            print(f'|{dest}| encoding failed!', file=sys.stderr, flush=True)
                        with open(dest, 'bw') as f:
                            f.write(buf)
                    print(f'|{dest}|', flush=True)
                    stat_ref['n_proc'] += 1
            print_stat()
    for xpath in xpaths:
        if os.path.isdir(xpath):
            _check_dir_pic_exts(xset_ref, xpath, stat_ref, xexts)


if '__main__' == __name__:
    import argparse

    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('target_dir', help='The target directory to check.')
    args = parser.parse_args()

    target_dir = args.target_dir
    check_dir_pic_exts(target_dir, set(['webp', 'jfif']))
